package cn.baddonkey.photography.agent.controller;

import cn.baddonkey.photography.agent.config.RedisCache;
import cn.baddonkey.photography.agent.core.PlatformResponse;
import cn.baddonkey.photography.agent.entity.dto.ChatCreateDto;
import cn.baddonkey.photography.agent.graph.GraphProcess;
import cn.baddonkey.photography.agent.utils.SecurityUtils;
import cn.hutool.core.util.IdUtil;
import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @author yingzi
 * @since 2025/6/13
 */
@Slf4j
@RestController
@RequestMapping("/graph/parallel-stream")
public class GraphStreamController {
    public static final ConcurrentHashMap<String, SseEmitter> SSE_POOL = new ConcurrentHashMap<>();
    private static final Logger logger = LoggerFactory.getLogger(GraphStreamController.class);

    private final CompiledGraph compiledGraph;

    @Autowired
    private RedisCache redisCache;

    public GraphStreamController(@Qualifier("parallelStreamGraph") StateGraph stateGraph) throws GraphStateException {
        this.compiledGraph = stateGraph.compile();
    }

//    @GetMapping(value = "/expand-translate", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expand(@RequestParam(value = "imageUrl") String imageUrl) throws GraphRunnerException {

        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(String.valueOf(Thread.currentThread().getId())).build();
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("imageUrl", imageUrl);

        GraphProcess graphProcess = new GraphProcess(this.compiledGraph);
        Sinks.Many<ServerSentEvent<String>> sink = Sinks.many().unicast().onBackpressureBuffer();
        AsyncGenerator<NodeOutput> resultFuture = compiledGraph.stream(objectMap, runnableConfig);
        graphProcess.processStream(resultFuture, sink);

        return sink.asFlux()
                .doOnCancel(() -> logger.info("Client disconnected from stream"))
                .doOnError(e -> logger.error("Error occurred during streaming", e));
    }

    @PostMapping("/chat/create")
    public PlatformResponse chatCreate(@RequestBody @Validated ChatCreateDto chatCreateDto) {
        String userId = SecurityUtils.getUserId();
        String uuid = IdUtil.simpleUUID();
        redisCache.setCacheObject("ChatCreate:" + uuid, chatCreateDto.getImageUrl(), 30, TimeUnit.SECONDS);
        return PlatformResponse.ok(uuid);
    }

    @GetMapping(value = "/expand-translate/v2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter expandv2(@RequestParam(value = "key") String key) throws GraphRunnerException {
        String imageUrl = redisCache.getCacheObject("ChatCreate:" + key);
        String threadId = String.valueOf(Thread.currentThread().getId());
        SseEmitter emitter = new SseEmitter(0L);
        SSE_POOL.put(threadId, emitter);
        // 👇 关键：先发送一个空注释事件（:keepalive\n\n），强制 flush 响应头
        try {
            emitter.send(SseEmitter.event().comment("connected"));
        } catch (IOException e) {
            emitter.completeWithError(e);
            throw new RuntimeException("Failed to initialize SSE", e);
        }

        emitter.onCompletion(() -> {
            log.info("onCompletion!");
            SSE_POOL.remove(threadId);
        });
        emitter.onError((e) -> {
            log.error("onError!", e);
            SSE_POOL.remove(threadId);
        });
        emitter.onTimeout(() -> {
            log.info("onTimeout!");
            SSE_POOL.remove(threadId);
        });
        Map<String, Object> objectMap = Map.of("imageUrl", imageUrl,
                "emitter", emitter,
                "threadId", threadId);



        CompletableFuture.runAsync(() -> {
            try {
                compiledGraph.call(objectMap);
                emitter.send("[DONE]");
                emitter.complete();
            } catch (Exception ex) {
                log.error("Graph execution error", ex);
                emitter.completeWithError(ex);
            }
        });
        System.out.println("================================");

        return emitter;
    }

}